import paddle
import json
import gzip
import numpy as np
from paddle.fluid.io import load_persistables
import paddle.nn.functional as F
from paddle.nn import Conv2D, MaxPool2D, Linear
import paddle.distributed as dist
import random
from PIL import Image
import argparse
import matplotlib.pyplot as plt


# 数据集相关参数，图片高度IMG_ROWS, 图片宽度IMG_COLS
IMG_ROWS = 28
IMG_COLS = 28

# 异步加载数据
ASYNC_LOAD = 0
# 读入数据时用到的batchsize
BATCHSIZE = 100
# 训练EPOCH数量
EPOCH_NUM = 4
# 使用GPU
USE_GPU = True
# 分布式训练
DIST_TRAIN = False
# 基于launch启动
# 使用本机全部GPU
# python -m paddle.distributed.launch train.py
# 使用指定GPU
# python -m paddle.distributed.launch --gpus '0,1' --log_dir ./mylog train.py
# 基于spawn启动
# # 启动train多进程训练，默认使用所有可见的GPU卡。
# if __name__ == '__main__':
#     dist.spawn(train)

# # 启动train函数2个进程训练，默认使用当前可见的前2张卡。
# if __name__ == '__main__':
#     dist.spawn(train, nprocs=2)

# # 启动train函数2个进程训练，默认使用第4号和第5号卡。
# if __name__ == '__main__':
#     dist.spawn(train, nprocs=2, selelcted_gpus='4,5')



# 异步加载数据
# 创建一个类MnistDataset，继承paddle.io.Dataset 这个类
# MnistDataset的作用和上面load_data()函数的作用相同，均是构建一个迭代器
class MnistDataset(paddle.io.Dataset):
    def __init__(self, mode):
        datafile = './work/mnist.json.gz'
        data = json.load(gzip.open(datafile))
        # 读取到的数据区分训练集，验证集，测试集
        train_set, val_set, eval_set = data
        if mode=='train':
            # 获得训练数据集
            imgs, labels = train_set[0], train_set[1]
            # print(len(imgs))
        elif mode=='valid':
            # 获得验证数据集
            imgs, labels = val_set[0], val_set[1]
        elif mode=='eval':
            # 获得测试数据集
            imgs, labels = eval_set[0], eval_set[1]
        else:
            raise Exception("mode can only be one of ['train', 'valid', 'eval']")
        
        # 校验数据
        imgs_length = len(imgs)
        assert len(imgs) == len(labels), \
            "length of train_imgs({}) should be the same as train_labels({})".format(len(imgs), len(labels))
        
        self.imgs = imgs
        self.labels = labels

    def __getitem__(self, idx):
        # 使用全连接层
        # img = np.array(self.imgs[idx]).astype('float32')
        # label = np.array(self.labels[idx]).astype('int64')
        # 使用卷积层
        img = np.reshape(self.imgs[idx], [1, IMG_ROWS, IMG_COLS]).astype('float32')
        label = np.reshape(self.labels[idx], [1]).astype('int64')
        
        return img, label

    def __len__(self):
        return len(self.imgs)

# # 声明数据加载函数，使用训练模式，MnistDataset构建的迭代器每次迭代只返回batch=1的数据
# train_dataset = MnistDataset(mode='train')
# # 使用paddle.io.DataLoader 定义DataLoader对象用于加载Python生成器产生的数据，
# # DataLoader 返回的是一个批次数据迭代器，并且是异步的；
# data_loader = paddle.io.DataLoader(train_dataset, batch_size=BATCHSIZE, shuffle=True)
# # 迭代的读取数据并打印数据的形状
# # for i, data in enumerate(data_loader()):
# #     images, labels = data
# #     print(i, images.shape, labels.shape)

# 同步加载数据
# 定义数据集读取器
def load_data(mode='train'):

    # 加载数据
    datafile = './work/mnist.json.gz'
    print('loading mnist dataset from {} ......'.format(datafile))
    data = json.load(gzip.open(datafile))
    print('mnist dataset load done')

    # 读取到的数据区分训练集，验证集，测试集
    train_set, val_set, eval_set = data


    if mode == 'train':
        # 获得训练数据集
        imgs, labels = train_set[0], train_set[1]
    elif mode == 'valid':
        # 获得验证数据集
        imgs, labels = val_set[0], val_set[1]
    elif mode == 'eval':
        # 获得测试数据集
        imgs, labels = eval_set[0], eval_set[1]
    else:
        raise Exception("mode can only be one of ['train', 'valid', 'eval']")

    #校验数据
    imgs_length = len(imgs)
    assert len(imgs) == len(labels), \
          "length of train_imgs({}) should be the same as train_labels({})".format(
                  len(imgs), len(labels))

    # 定义数据集每个数据的序号， 根据序号读取数据
    index_list = list(range(imgs_length))

    
    # 定义数据生成器
    def data_generator():
        if mode == 'train':
            random.shuffle(index_list)
        imgs_list = []
        labels_list = []
        for i in index_list:
            # img = np.array(imgs[i]).astype('float32')
            # print('img.shape:', img.shape)
            # label = np.array(labels[i]).astype('int64')
            # 在使用卷积神经网络结构时，uncomment 下面两行代码
            img = np.reshape(imgs[i], [1, IMG_ROWS, IMG_COLS]).astype('float32')
            label = np.reshape(labels[i], [1]).astype('int64')
            imgs_list.append(img) 
            labels_list.append(label)
            if len(imgs_list) == BATCHSIZE:
                yield np.array(imgs_list), np.array(labels_list)
                imgs_list = []
                labels_list = []

        # 如果剩余数据的数目小于BATCHSIZE，
        # 则剩余数据一起构成一个大小为len(imgs_list)的mini-batch
        if len(imgs_list) > 0:
            yield np.array(imgs_list), np.array(labels_list)

    return data_generator

# 多层卷积神经网络实现
class MNIST(paddle.nn.Layer):
    def __init__(self):
        super(MNIST, self).__init__()
       
        # 定义卷积层，输出特征通道out_channels设置为20，卷积核的大小kernel_size为5，卷积步长stride=1，padding=2
        self.conv1 = Conv2D(in_channels=1, out_channels=20, kernel_size=5, stride=1, padding=2)
        # 定义池化层，池化核的大小kernel_size为2，池化步长为2
        self.max_pool1 = MaxPool2D(kernel_size=2, stride=2)
        # 定义卷积层，输出特征通道out_channels设置为20，卷积核的大小kernel_size为5，卷积步长stride=1，padding=2
        self.conv2 = Conv2D(in_channels=20, out_channels=20, kernel_size=5, stride=1, padding=2)
        # 定义池化层，池化核的大小kernel_size为2，池化步长为2
        self.max_pool2 = MaxPool2D(kernel_size=2, stride=2)
        # 定义一层全连接层，输出维度是10
        self.fc = Linear(in_features=980, out_features=10)
         
    # 定义网络前向计算过程，卷积后紧接着使用池化层，最后使用全连接层计算最终输出
    # 卷积层激活函数使用Relu，全连接层不使用激活函数
    def forward(self, inputs, label, check_shape=False, check_content=False):
        x1 = self.conv1(inputs)
        x2 = F.relu(x1)
        x3 = self.max_pool1(x2)
        x4 = self.conv2(x3)
        x5 = F.relu(x4)
        x6 = self.max_pool2(x5)
        x6 = paddle.reshape(x6, [x6.shape[0], -1])
        x7 = self.fc(x6)
       #  x = F.softmax(x)
        if label is None:
            return x7
        else:
            acc = paddle.metric.accuracy(input=x7, label=label)
            return x7, acc
        if check_shape:
            # 打印每层网络设置的超参数-卷积核尺寸，卷积步长，卷积padding，池化核尺寸
            print("\n########## print network layer's superparams ##############")
            print("conv1-- kernel_size:{}, padding:{}, stride:{}".format(self.conv1.weight.shape, self.conv1._padding, self.conv1._stride))
            print("conv2-- kernel_size:{}, padding:{}, stride:{}".format(self.conv2.weight.shape, self.conv2._padding, self.conv2._stride))
            #print("max_pool1-- kernel_size:{}, padding:{}, stride:{}".format(self.max_pool1.pool_size, self.max_pool1.pool_stride, self.max_pool1._stride))
            #print("max_pool2-- kernel_size:{}, padding:{}, stride:{}".format(self.max_pool2.weight.shape, self.max_pool2._padding, self.max_pool2._stride))
            print("fc-- weight_size:{}, bias_size_{}".format(self.fc.weight.shape, self.fc.bias.shape))
            
            # 打印每层的输出尺寸
            print("\n########## print shape of features of every layer ###############")
            print("inputs_shape: {}".format(inputs.shape))
            print("outputs1_shape: {}".format(x1.shape))
            print("outputs2_shape: {}".format(x2.shape))
            print("outputs3_shape: {}".format(x3.shape))
            print("outputs4_shape: {}".format(x4.shape))
            print("outputs5_shape: {}".format(x5.shape))
            print("outputs6_shape: {}".format(x6.shape))
            print("outputs7_shape: {}".format(x7.shape))
        if check_content:
            # 打印卷积层的参数-卷积核权重，权重参数较多，此处只打印部分参数
            print("\n########## print convolution layer's kernel ###############")
            print("conv1 params -- kernel weights:", self.conv1.weight[0][0])
            print("conv2 params -- kernel weights:", self.conv2.weight[0][0])

            # 创建随机数，随机打印某一个通道的输出值
            idx1 = np.random.randint(0, x1.shape[1])
            idx2 = np.random.randint(0, x4.shape[1])
            # 打印卷积-池化后的结果，仅打印batch中第一个图像对应的特征
            print("\nThe {}th channel of conv1 layer: ".format(idx1), x1[0][idx1])
            print("The {}th channel of conv2 layer: ".format(idx2), x4[0][idx2])
            print("The output of last layer:", x7[0], '\n')


def train(model):
    if DIST_TRAIN:
        dist.init_parallel_env()
        model = paddle.DataParallel(model)
    paddle.set_device('gpu:0') if USE_GPU else paddle.set_device('cpu')
    model.train()
    if ASYNC_LOAD:
        train_dataset = MnistDataset(mode='train')
        train_loader = paddle.io.DataLoader(train_dataset, batch_size=BATCHSIZE, shuffle=True)
    else:
        train_loader = load_data('train')
    #四种优化算法的设置方案，可以逐一尝试效果
    # opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
    # opt = paddle.optimizer.Momentum(learning_rate=0.01, momentum=0.9, parameters=model.parameters())
    # opt = paddle.optimizer.Adagrad(learning_rate=0.01, parameters=model.parameters())
    opt = paddle.optimizer.Adam(learning_rate=0.01, weight_decay=paddle.regularizer.L2Decay(coeff=1e-5), parameters=model.parameters())

    # 加载参数到模型
    # params_dict = paddle.load('./checkpoint/mnist_epoch0.pdparams')
    # opt_dict = paddle.load('./checkpoint/mnist_epoch0.pdopt')
    # model.set_state_dict(params_dict)
    # opt.set_state_dict(opt_dict)
    
    iter = 0
    iters=[]
    losses=[]
    for epoch_id in range(EPOCH_NUM):
        for batch_id, data in enumerate(train_loader()):
            images, labels = data
            images = paddle.to_tensor(images)
            images = paddle.reshape(images, [images.shape[0], 1, IMG_ROWS, IMG_COLS])
            labels = paddle.to_tensor(labels)
            
            #前向计算的过程  
            predicts, acc = model(images, labels)

            #计算损失，取一个批次样本损失的平均值
            loss = F.cross_entropy(predicts, labels)
            avg_loss = paddle.mean(loss)
            
            #每训练了200批次的数据，打印下当前Loss的情况
            if batch_id % 200 == 0:
                print("epoch: {}, batch: {}, loss is: {}, acc is: {}".format(epoch_id, batch_id, avg_loss.numpy(), acc.numpy()))
                iters.append(iter)
                losses.append(avg_loss.numpy())
                iter += 200

            #后向传播，更新参数的过程
            avg_loss.backward()
            opt.step()
            opt.clear_grad()

    # 保存checkpoints
    # paddle.save(opt.state_dict(), './{}/mnist_epoch{}'.format(save_path,i)+'.pdopt')
    # paddle.save(model.state_dict(), './{}/mnist_epoch{}'.format(save_path,i)+'.pdparams')
    # 保存模型参数
    paddle.save(model.state_dict(), 'mnist.pdparams')
    print('Model has been saved.')
    return iters, losses


# 读取一张本地的样例图片，转变成模型输入的格式
def load_image(img_path):
    # 从img_path中读取图像，并转为灰度图
    im = Image.open(img_path).convert('L')
    im = im.resize((28, 28), Image.ANTIALIAS)
    im = np.array(im).reshape(1, 1, 28, 28).astype(np.float32)
    # 图像归一化
    im = 1.0 - im / 255.
    return im


def evaluation(model):
    print('start evaluation .......')
    # 定义预测过程
    params_file_path = 'mnist.pdparams'
    # 加载模型参数
    param_dict = paddle.load(params_file_path)
    model.load_dict(param_dict)

    model.eval()
    eval_loader = load_data('eval')

    acc_set = []
    avg_loss_set = []
    for batch_id, data in enumerate(eval_loader()):
        images, labels = data
        images = paddle.to_tensor(images)
        labels = paddle.to_tensor(labels)
        predicts, acc = model(images, labels)
        loss = F.cross_entropy(input=predicts, label=labels)
        avg_loss = paddle.mean(loss)
        acc_set.append(float(acc.numpy()))
        avg_loss_set.append(float(avg_loss.numpy()))
    
    #计算多个batch的平均损失和准确率
    acc_val_mean = np.array(acc_set).mean()
    avg_loss_val_mean = np.array(avg_loss_set).mean()

    print('loss={}, acc={}'.format(avg_loss_val_mean, acc_val_mean))


def create_parser():
    # 实例化
    # 创建一个 ArgumentParser 对象
    # ArgumentParser 对象包含将命令行解析成 Python 数据类型所需的全部信息。
    parser = argparse.ArgumentParser(description = 'Handwriting number recognition program.')

    parser.add_argument('-t', '--train', help='Start train process.', action="store_true")
    parser.add_argument('-p', '--predict', help='Predict numbers from picture.', action="store_true")
    parser.add_argument('-e', '--evaluate', help='Evaluate model from dataset.', action="store_true")
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = create_parser()
    if args.train:
        #创建模型
        model = MNIST()
        #启动训练过程
        iters, losses = train(model)
        plt.figure()
        plt.title("train loss", fontsize=24)
        plt.xlabel("iter", fontsize=14)
        plt.ylabel("loss", fontsize=14)
        plt.plot(iters, losses,color='red',label='train loss') 
        plt.grid()
        plt.show()
    elif args.evaluate:
        model = MNIST()
        evaluation(model)
    elif args.predict:
        # 定义预测过程
        model = MNIST()
        params_file_path = 'mnist.pdparams'
        img_path = 'work/example_5.jpg'
        # 加载模型参数
        param_dict = paddle.load(params_file_path)
        model.load_dict(param_dict)
        # 灌入数据
        model.eval()
        tensor_img = load_image(img_path)
        # print(paddle.to_tensor(tensor_img).shape)
        #模型反馈10个分类标签的对应概率
        results = model(paddle.to_tensor(tensor_img), label=None)
        #取概率最大的标签作为预测输出
        lab = np.argsort(results.numpy())
        print("本次预测的数字是: ", lab[0][-1])

